/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG//内置宏ASCENDC_CPU_DEBUG是区分运行CPU模式或NPU模式逻辑的标志
#include "acl/acl.h"
extern void is_real_do(uint32_t blockDim, void *stream, 
                                void* x, void* z, int32_t dtype);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void is_real_custom(GM_ADDR x, GM_ADDR z, int32_t dtype);
#endif

int32_t main(int32_t argc, char* argv[])
{
    int32_t dtype = 1; 
    uint32_t blockDim = 8;
    size_t inputByteSize = 8 * 2048 * sizeof(uint16_t);
    size_t outputByteSize = 8 * 2048 * sizeof(int8_t); 
    



#ifdef ASCENDC_CPU_DEBUG
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *z = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    ReadFile("./input/input_x.bin", inputByteSize, x, inputByteSize);
    printf("Read input data, size: %zu bytes\n", inputByteSize);
    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(is_real_custom, blockDim, x, z, dtype);

    WriteFile("./output/output_z.bin", z, outputByteSize);
    printf("Writing output data, size: %zu bytes\n", outputByteSize);
    AscendC::GmFree(x);
    AscendC::GmFree(z);

#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *xHost, *zHost;
    uint8_t *xDevice, *zDevice;

    CHECK_ACL(aclrtMallocHost((void**)&xHost, inputByteSize));
    CHECK_ACL(aclrtMallocHost((void**)&zHost, outputByteSize));
    CHECK_ACL(aclrtMalloc((void**)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void**)&zDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_x.bin", inputByteSize, xHost, inputByteSize);

    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    // 调用 kernel
    is_real_do(blockDim, stream, xDevice, zDevice, dtype);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    // 拷贝回 host
    CHECK_ACL(aclrtMemcpy(zHost, outputByteSize, zDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));

    WriteFile("./output/output_z.bin", zHost, outputByteSize);
    // 释放
    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(zHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(0));
    CHECK_ACL(aclFinalize());
#endif

    return 0;
}
